# 解决OpenMP库冲突问题
import os

import agent

os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from scipy.ndimage import zoom

from agent import Agent, AGENT_PARAMS
from target_probability_map import TargetProbabilityMap

from concurrent.futures import ThreadPoolExecutor, as_completed


class MultiTargetProbabilityMap:
    def __init__(self, width, height, grid_resolution=1000, kernel_size=7, sigma=1.0):
        """
        多目标概率地图管理器

        参数:
        width, height: 实际战场尺寸（米）
        grid_resolution: 网格分辨率（米/网格单元）
        kernel_size: 扩散核大小
        sigma: 高斯扩散核标准差
        """
        self.width = width
        self.height = height
        self.grid_resolution = grid_resolution

        # 计算网格尺寸
        self.grid_width = int(np.ceil(width / grid_resolution))
        self.grid_height = int(np.ceil(height / grid_resolution))

        # 目标概率地图字典 {target_id: TargetProbabilityMap}
        self.target_maps = {}
        # 不可通行区域障碍地图
        self.obstacle_map = np.zeros((self.grid_width, self.grid_height))

        # 默认地图参数
        self.kernel_size = kernel_size
        self.sigma = sigma

    def add_target(self, target_id: int, initial_pos=None, max_speed: float = agent.AGENT_PARAMS['BSV_MAX_VEL'],
                   time: float = 0.0):
        """
        添加新目标

        参数:
        target_id: 目标唯一标识符
        initial_pos: 初始位置 (x, y) 实际坐标
        """
        if target_id in self.target_maps:
            if initial_pos is not None:
                self.target_maps[target_id].update_position(self.real_to_grid(initial_pos), [0, 0], 0.0)
            return

        # 创建初始概率分布
        initial_prob = None
        if initial_pos is not None:
            grid_x, grid_y = self.real_to_grid(initial_pos)
            initial_prob = np.zeros((self.grid_width, self.grid_height))
            if 0 <= grid_x < self.grid_width and 0 <= grid_y < self.grid_height:
                initial_prob[grid_x, grid_y] = 1.0

        # 创建目标地图
        self.target_maps[target_id] = TargetProbabilityMap(
            self.grid_resolution,
            target_id,
            self.grid_width,
            self.grid_height,
            initial_prob,
            max_speed,
            time
        )

    def add_target_with_prob(self, target_id: int, initial_prob: np.ndarray,
                             max_speed: float = agent.AGENT_PARAMS['BSV_MAX_VEL'],
                             time: float = 0.0):
        """
        添加新目标

        参数:
        target_id: 目标唯一标识符
        initial_prob: 初始概率分布 (二维 numpy 数组)
        """

        # 计算缩放因子
        zoom_factor_width = self.grid_width / initial_prob.shape[0]
        zoom_factor_height = self.grid_height / initial_prob.shape[1]

        # 缩放初始概率分布
        initial_prob = zoom(initial_prob, zoom=(zoom_factor_width, zoom_factor_height), order=1)

        # 创建目标地图
        self.target_maps[target_id] = TargetProbabilityMap(
            self.grid_resolution,
            target_id,
            self.grid_width,
            self.grid_height,
            initial_prob,
            max_speed,
            time
        )

    def remove_target(self, target_id: int):
        """移除目标"""
        if target_id in self.target_maps:
            self.target_maps[target_id].belief = 0.0
            self.target_maps[target_id].prob_map = np.zeros((self.grid_width, self.grid_height))

    def real_to_grid(self, real_pos):
        """实际坐标转网格坐标"""
        x, y = real_pos
        grid_x = int(x / self.grid_resolution)
        grid_y = int(y / self.grid_resolution)
        return grid_x, grid_y

    def predict(self):
        """
        对所有目标执行预测步（多线程版本）
        """
        with ThreadPoolExecutor() as executor:
            # 提交所有目标地图的预测任务
            futures = []
            for target_map in self.target_maps.values():
                future = executor.submit(target_map.predict)
                futures.append(future)

            # 等待所有任务完成并处理可能的异常
            for future in as_completed(futures):
                try:
                    future.result()
                except Exception as e:
                    print(f"Error during prediction: {e}")

    def update(self, detector_agent: Agent, detected_targets: ['Agent']):
        """
        根据探测器观测更新所有目标地图（多线程版本）

        参数:
        detector_agent: 进行探测的智能体
        detected_targets: 探测到的目标列表 (Agent对象)
        """
        detected_ids = {target.id for target in detected_targets}

        # 添加新发现的目标
        for target_id in detected_ids:
            if target_id not in self.target_maps:
                self.add_target(target_id)

        # 使用线程池并行更新所有目标地图
        with ThreadPoolExecutor() as executor:
            # 提交任务到线程池
            futures = []
            for target_id, target_map in self.target_maps.items():
                future = executor.submit(self._update_single_target,
                                         target_id, target_map, detected_ids,
                                         detected_targets, detector_agent)
                futures.append(future)

            # 等待所有任务完成
            for future in as_completed(futures):
                try:
                    future.result()
                except Exception as e:
                    print(f"Error updating target map: {e}")

    def _update_single_target(self, target_id, target_map, detected_ids, detected_targets, detector_agent):
        """
        更新单个目标地图的辅助方法
        """
        if target_id in detected_ids:
            # 如果目标被探测到，更新其精确位置
            target = next(t for t in detected_targets if t.id == target_id)
            target_pos = (target.x, target.y)
            target_map.update_position(target_pos, [target.vx, target.vy], target.world.time)
        else:
            # 如果目标未被探测到，更新其不存在概率
            target_map.update_absence(detector_agent, detector_agent.world.time)

    def update_entropy(self):
        """
        多线程更新所有目标地图的信息熵
        """
        with ThreadPoolExecutor() as executor:
            # 提交所有目标地图的熵计算任务
            futures = []
            for target_map in self.target_maps.values():
                future = executor.submit(target_map.compute_entropy)
                futures.append(future)

            # 等待所有任务完成并处理可能的异常
            for future in as_completed(futures):
                try:
                    future.result()
                except Exception as e:
                    print(f"Error computing entropy: {e}")

    def get_target_map(self, target_id):
        """获取特定目标的概率地图"""
        return self.target_maps.get(target_id)

    def get_all_maps(self):
        """获取所有目标的概率地图"""
        return self.target_maps

    def get_obstacle_map(self,
                         cumulative_prob_threshold=0.8,
                         dilation_radius=AGENT_PARAMS['BSV_SCAN_RANGE'] + 1000,
                         min_belief=0.1) -> np.ndarray:
        """
        生成不可通行区域地图（优化膨胀操作性能）

        参数:
        cumulative_prob_threshold: 累积概率阈值 (0-1)
        dilation_radius: 不可通行区域的扩展半径 (实际尺度)
        min_belief: 目标存在的最小概率阈值，低于此值的目标不参与计算

        返回:
        obstacle_map: 二维布尔数组，True表示该网格不可通行
        """
        # 实际半径转网格单元
        grid_dilation_radius = int(dilation_radius / self.grid_resolution)
        # 初始化障碍地图
        obstacle_map = np.zeros((self.grid_width, self.grid_height), dtype=bool)

        # 分离已发现目标和未发现目标
        detected_targets = []
        undetected_targets = []

        for target_map in self.target_maps.values():
            if target_map.belief < min_belief:
                continue
            # 已发现目标：当前时间与最后发现时间相同
            if target_map.last_detected == target_map.time:
                detected_targets.append(target_map)
            else:
                undetected_targets.append(target_map)

        # 1. 处理未发现目标 - 使用局部膨胀
        if undetected_targets:
            # 创建组合概率图
            combined_prob = np.zeros((self.grid_width, self.grid_height))
            total_belief = 0.0
            for target_map in undetected_targets:
                # 只处理概率大于阈值的部分
                if np.any(target_map.prob_map > 1e-6):  # 微小概率阈值
                    combined_prob += target_map.prob_map * target_map.belief
                    total_belief += target_map.belief
            if total_belief > 0:
                # 归一化
                combined_prob /= combined_prob.sum()

                # 优化排序：只处理非零点
                non_zero_indices = np.nonzero(combined_prob > 0)
                non_zero_probs = combined_prob[non_zero_indices]

                if non_zero_probs.size > 0:
                    # 按概率降序排序非零点
                    sorted_idx = np.argsort(non_zero_probs)[::-1]
                    sorted_probs = non_zero_probs[sorted_idx]

                    # 累积概率直到达到阈值
                    cum_prob = 0.0
                    danger_points = []  # 存储危险点坐标
                    for i, prob in enumerate(sorted_probs):
                        cum_prob += prob
                        # 收集原始索引
                        x = non_zero_indices[0][sorted_idx[i]]
                        y = non_zero_indices[1][sorted_idx[i]]
                        danger_points.append((x, y))
                        if cum_prob >= cumulative_prob_threshold:
                            break

                    # 计算危险区域的最小包围盒
                    if danger_points:
                        xs, ys = zip(*danger_points)
                        min_x, max_x = min(xs), max(xs)
                        min_y, max_y = min(ys), max(ys)

                        # 扩展包围盒以包含膨胀区域
                        pad = grid_dilation_radius
                        roi_min_x = max(0, min_x - pad)
                        roi_max_x = min(self.grid_width, max_x + pad + 1)
                        roi_min_y = max(0, min_y - pad)
                        roi_max_y = min(self.grid_height, max_y + pad + 1)

                        # 创建ROI区域内的危险区地图
                        roi_width = roi_max_x - roi_min_x
                        roi_height = roi_max_y - roi_min_y
                        roi_danger = np.zeros((roi_width, roi_height), dtype=bool)

                        # 标记危险点（相对ROI坐标）
                        for x, y in danger_points:
                            roi_x = x - roi_min_x
                            roi_y = y - roi_min_y
                            if 0 <= roi_x < roi_width and 0 <= roi_y < roi_height:
                                roi_danger[roi_x, roi_y] = True

                        # 使用形态学膨胀
                        if grid_dilation_radius > 0:
                            # 创建圆形核
                            y, x = np.ogrid[-grid_dilation_radius:grid_dilation_radius + 1,
                                   -grid_dilation_radius:grid_dilation_radius + 1]
                            kernel = (x * x + y * y <= grid_dilation_radius * grid_dilation_radius)

                            # 对ROI区域进行膨胀
                            from scipy.ndimage import binary_dilation
                            roi_obstacle = binary_dilation(roi_danger, structure=kernel)

                            # 将ROI结果合并到完整地图
                            obstacle_map[roi_min_x:roi_max_x, roi_min_y:roi_max_y] = np.logical_or(
                                obstacle_map[roi_min_x:roi_max_x, roi_min_y:roi_max_y],
                                roi_obstacle
                            )
                        else:
                            # 直接标记危险点
                            for x, y in danger_points:
                                obstacle_map[x, y] = True

        # 2. 处理已发现目标 - 使用局部膨胀
        if detected_targets:
            # 创建点图标记目标位置
            point_map = np.zeros((self.grid_width, self.grid_height), dtype=bool)
            target_positions = []

            for target_map in detected_targets:
                grid_x, grid_y = self.real_to_grid(target_map.position)
                if 0 <= grid_x < self.grid_width and 0 <= grid_y < self.grid_height:
                    point_map[grid_x, grid_y] = True
                    target_positions.append((grid_x, grid_y))

            # 计算包围盒
            if target_positions:
                xs, ys = zip(*target_positions)
                min_x, max_x = min(xs), max(xs)
                min_y, max_y = min(ys), max(ys)

                # 扩展包围盒以包含膨胀区域
                pad = grid_dilation_radius
                roi_min_x = max(0, min_x - pad)
                roi_max_x = min(self.grid_width, max_x + pad + 1)
                roi_min_y = max(0, min_y - pad)
                roi_max_y = min(self.grid_height, max_y + pad + 1)

                # 创建ROI区域内的点图
                roi_width = roi_max_x - roi_min_x
                roi_height = roi_max_y - roi_min_y
                roi_points = np.zeros((roi_width, roi_height), dtype=bool)

                # 标记目标点（相对ROI坐标）
                for x, y in target_positions:
                    roi_x = x - roi_min_x
                    roi_y = y - roi_min_y
                    if 0 <= roi_x < roi_width and 0 <= roi_y < roi_height:
                        roi_points[roi_x, roi_y] = True

                # 使用形态学膨胀
                if grid_dilation_radius > 0:
                    # 创建圆形核
                    y, x = np.ogrid[-grid_dilation_radius:grid_dilation_radius + 1,
                           -grid_dilation_radius:grid_dilation_radius + 1]
                    kernel = (x * x + y * y <= grid_dilation_radius * grid_dilation_radius)

                    # 对ROI区域进行膨胀
                    from scipy.ndimage import binary_dilation
                    roi_obstacle = binary_dilation(roi_points, structure=kernel)

                    # 将ROI结果合并到完整地图
                    obstacle_map[roi_min_x:roi_max_x, roi_min_y:roi_max_y] = np.logical_or(
                        obstacle_map[roi_min_x:roi_max_x, roi_min_y:roi_max_y],
                        roi_obstacle
                    )
                else:
                    # 直接标记目标点
                    for x, y in target_positions:
                        obstacle_map[x, y] = True

        self.obstacle_map = obstacle_map
        return obstacle_map

    def can_pass(self, position) -> bool:
        """
        判断一个位置是否安全可通行

        :param position: 二维数组(真实尺度)
        :return: True: 安全可通行 False: 存在敌方打击风险
        """
        try:
            # 转换为网格坐标
            grid_x, grid_y = self.real_to_grid(position)

            # 检查是否在地图范围内
            if not (0 <= grid_x < self.grid_width and 0 <= grid_y < self.grid_height):
                return False  # 超出地图边界视为不可通行

            # 检查障碍地图 - 直接检查布尔值
            return not self.obstacle_map[grid_x, grid_y]
        except Exception as e:
            # 处理可能的异常（如位置格式错误）
            import traceback
            print("=" * 80)
            print("ERROR in MultiTargetProbabilityMap.can_pass()")
            print("=" * 80)
            traceback.print_exc(limit=None)  # 打印完整调用栈
            print("=" * 80)
            print(f"Error details: {e}")
            print(f"Position parameter: {position}")
            print(f"Position type: {type(position)}")
            if hasattr(self, 'grid_width'):
                print(f"Grid dimensions: {self.grid_width} x {self.grid_height}")
            print("=" * 80)
            return False

    def plot(self, ax, title="Multi-Target Probability Map", cmap='viridis'):
        """可视化所有目标的概率地图（合并显示）"""
        combined_map = np.zeros((self.grid_width, self.grid_height))

        # self.get_obstacle_map()
        # 合并所有目标的概率
        for target_map in self.target_maps.values():
            # 缩放熵值过小的目标的显示值以更好地显示结果
            combined_map += target_map.prob_map * target_map.belief

        # 归一化
        if combined_map.sum() > 0:
            combined_map = np.log1p(
                self.width * self.height * combined_map / self.grid_resolution / self.grid_resolution)  # 对数变换增强低值区域
            # combined_map /= combined_map.sum()  # 归一化

        # 可视化
        ax.clear()
        ax.imshow(combined_map.transpose(), cmap=cmap, origin='lower',
                  extent=[0, self.width, 0, self.height])
        ax.set_title(title)
        ax.set_xlabel('X坐标 (m)')
        ax.set_ylabel('Y坐标 (m)')
        ax.grid(False)

        # 添加目标ID标注
        for target_id, target_map in self.target_maps.items():
            real_x = target_map.position[0]
            real_y = target_map.position[1]

            # 添加标注
            ax.text(real_x, real_y, f'T{target_id}\nB:{target_map.belief:.2f}',
                    fontsize=8, ha='center', va='center', color='white',
                    bbox=dict(facecolor='black', alpha=0.1))

    def __getitem__(self, item: int) -> TargetProbabilityMap:
        """
        获取指定id的单层态势图
        :param item: 目标id，int
        :return: 指定id的单层态势图
        """
        return self.target_maps[item]

    def __len__(self):
        """
        获取MultiTargetProbabilityMap实例中包含的TargetProbabilityMap实例的数量
        :return: 目标数量
        """
        return len(self.target_maps)


# 示例用法
if __name__ == "__main__":
    # 模拟战场尺寸 100km x 100km，网格分辨率1km
    world_width = 100000
    world_height = 100000
    grid_resolution = 1000

    # 创建多目标概率地图
    multi_map = MultiTargetProbabilityMap(world_width, world_height, grid_resolution)

    # 添加3个目标
    multi_map.add_target(1, (20000, 50000))
    multi_map.add_target(2, (60000, 30000))
    multi_map.add_target(3, (40000, 70000))

    # 创建图形
    fig, ax = plt.subplots(figsize=(10, 8))

    # 初始状态可视化
    multi_map.plot(ax, title="Initial Multi-Target Probability Map")


    # 模拟探测器类
    class Detector:
        def __init__(self, x, y, sensor_range):
            self.x = x
            self.y = y
            self.sensor_range = sensor_range


    # 模拟目标类
    class Target:
        def __init__(self, id, x, y):
            self.id = id
            self.x = x
            self.y = y


    def update_frame(i, multi_map, ax):
        # 执行预测步
        multi_map.predict()

        # 创建探测器（位置随时间变化）
        detector = Detector(
            x=30000 + i * 5000,
            y=40000 + i * 3000,
            sensor_range=20000
        )

        # 模拟探测结果（随机探测到部分目标）
        detected_targets = []
        if i % 3 == 0:
            detected_targets.append(Target(1, 20000 + i * 1000, 50000))
        if i % 4 == 0:
            detected_targets.append(Target(2, 60000, 30000 + i * 500))
        if i % 5 == 0:
            detected_targets.append(Target(3, 40000 - i * 300, 70000))

        # 更新概率地图
        multi_map.update(detector, detected_targets)

        # 可视化
        multi_map.plot(ax, title=f"Step {i + 1} - Multi-Target Probability Map")


    # 创建动画
    ani = FuncAnimation(fig, update_frame, fargs=(multi_map, ax),
                        frames=20, repeat=False, interval=500)

    plt.show()
